#include <tinx/cpu.h>
#include <tinx/task.h>
#include <tinx/debug.h>
#include <tinx/global.h>
#include <tinx/assert.h>
#include <tinx/arena.h>
#include <tinx/memory.h>
#include <tinx/string.h>

bool fpu_check()
{
    cpu_version_t ver;
    cpu_version(&ver);
    if (!ver.FPU)
        return false;

    u16 testword = 0xAA55;

    asm volatile(
        "movl %%cr0, %%eax\n"
        "andl %0, %%eax\n"
        "movl %%eax, %%cr0\n"
        "fninit\n"
        "fwait\n"
        "fnstsw %1\n" :: "g"(~(CR0_TS | CR0_EM)), "m"(testword) : "eax", "memory");

    return testword == 0;
}

void fpu_sse_handler()
{
    static task_t *last_task = NULL;
    task_t *task = running_task();

    set_cr0(get_cr0() & ~CR0_TS);

    if (task == last_task)
        return;

    if (last_task && task->flags & TASK_FLAG_FPU_SSE_USED)
    {
        asm volatile("fxsave %0\n" :: "m"(last_task->fpusse) : "memory");
        task->flags |= TASK_FLAG_FPU_SSE_USED;
    }

    last_task = task;

    if (task->fpusse)
    {
        asm volatile("fxrstor %0\n" :: "m"(task->fpusse) : "memory");
    }
    else
    {
        task->fpusse = (u8 *)alloc_kpage(1);
        memset(task->fpusse, 0, PAGE_SIZE);

        asm volatile(
            "fnclex\n"
            "fninit\n");

        task->flags |= TASK_FLAG_FPU_SSE_USED;
    }
}

bool sse_check()
{
    cpu_version_t ver;
    cpu_version(&ver);

    if (!ver.SSE)
        return false;
    if (!ver.SSE2)
        return false;
    if (!ver.FXSR)
        return false;

    return true;
}

void sse_enable()
{
    u32 mxcsr = 0x1f08;
    asm volatile(
        "movl %%cr0, %%eax\n"
        "andl %0, %%eax\n"
        "orl %1, %%eax\n"
        "movl %%eax, %%cr0\n"
        "movl %%cr4, %%eax\n"
        "orl %2, %%eax\n"
        "movl %%eax, %%cr4\n"
        "movl %3, %%eax\n"
        "ldmxcsr (%%eax)\n"
        :: "g"(~CR0_EM), "g"(CR0_MP), "g"(0b11 << 9), "g"(&mxcsr) : "eax");
}

void fpu_sse_init()
{
    assert(fpu_check());
    assert(sse_check());

    sse_enable();

    set_cr0(get_cr0() | CR0_TS);

    set_interrupt_handler(INTR_NM, fpu_sse_handler);
}
